import heapq
from collections import defaultdict
from itertools import product
import matplotlib.pyplot as plt
import numpy as np


def space_time_a_star(start_positions, goal_positions, map, num_agents):
    # 将起始位置转换为元组
    start_positions = tuple(start_positions)
    goal_positions = tuple(goal_positions)

    # 初始化
    open_list = []
    closed_set = set()
    g_scores = defaultdict(lambda: float('inf'))
    f_scores = defaultdict(lambda: float('inf'))
    came_from = {}  # 用于存储路径信息

    # 初始状态
    initial_state = (start_positions, 0)  # (positions, time_step)
    g_scores[initial_state] = 0
    f_scores[initial_state] = heuristic(start_positions, goal_positions)

    # 将初始状态放入开放列表
    heapq.heappush(open_list, (f_scores[initial_state], initial_state))

    while open_list:
        _, current_state = heapq.heappop(open_list)
        current_positions, current_time = current_state

        if current_positions == goal_positions:
            return reconstruct_path(came_from, current_state)

        closed_set.add(current_state)

        for neighbor in get_neighbors(current_positions, current_time, map, num_agents):
            neighbor_positions, neighbor_time = neighbor

            if neighbor in closed_set:
                continue

            tentative_g_score = g_scores[current_state] + 1  # 移动一步的成本为1

            if tentative_g_score < g_scores[neighbor]:
                came_from[neighbor] = current_state  # 记录父节点
                g_scores[neighbor] = tentative_g_score
                f_scores[neighbor] = tentative_g_score + heuristic(neighbor_positions, goal_positions)
                heapq.heappush(open_list, (f_scores[neighbor], neighbor))

    return None  # 未找到解


def heuristic(positions, goals):
    # 返回启发式值，例如曼哈顿距离之和
    return sum(abs(p[0] - g[0]) + abs(p[1] - g[1]) for p, g in zip(positions, goals))


def get_neighbors(positions, current_time, map, num_agents):
    # 生成邻居状态，同时检查是否有冲突
    neighbors = []
    for i in range(num_agents):
        for dx, dy in [(-1, 0), (1, 0), (0, -1), (0, 1)]:
            new_x, new_y = positions[i][0] + dx, positions[i][1] + dy
            if is_valid(new_x, new_y, map):
                if has_conflict((new_x, new_y), positions, i, current_time):
                    print(f"Agent {i + 1} at time {current_time} cannot move to ({new_x}, {new_y}) due to conflict.")
                else:
                    new_positions = list(positions)
                    new_positions[i] = (new_x, new_y)
                    neighbors.append((tuple(new_positions), current_time + 1))
                    print(f"Agent {i + 1} at time {current_time} can move to ({new_x}, {new_y}).")
            else:
                print(
                    f"Agent {i + 1} at time {current_time} cannot move to ({new_x}, {new_y}) due to obstacle or out of bounds.")
    return neighbors


def is_valid(x, y, map):
    # 检查位置 (x, y) 是否在地图范围内且不是障碍物
    return 0 <= x < len(map) and 0 <= y < len(map[0]) and map[x][y] != 1


def has_conflict(new_position, positions, agent_index, current_time):
    # 检查新位置是否与其他智能体发生冲突
    for i, pos in enumerate(positions):
        if i != agent_index and new_position == pos:
            return True
    return False


def reconstruct_path(came_from, final_state):
    # 重构路径
    path = []
    current = final_state
    while current in came_from:
        path.append(current)
        current = came_from[current]
    path.append(current)
    path.reverse()  # 反转路径以获得从开始到结束的顺序
    return path
def visualize_path(path, map):
    # 获取地图尺寸
    height = len(map)
    width = len(map[0])

    # 创建一个新的图形
    fig, ax = plt.subplots()

    # 绘制栅格
    for x in range(width + 1):
        ax.axvline(x, color='black', linewidth=1)
    for y in range(height + 1):
        ax.axhline(y, color='black', linewidth=1)

    # 设置坐标轴范围
    ax.set_xlim(0, width)
    ax.set_ylim(0, height)

    # 反转y轴，使得原点在左下角
    ax.invert_yaxis()

    # 移除坐标轴标签
    ax.set_xticklabels([])
    ax.set_yticklabels([])

    # 显示障碍物
    for x in range(width):
        for y in range(height):
            if map[y][x] == 1:
                rect = plt.Rectangle((x, y), 2, 1, facecolor='gray')
                ax.add_patch(rect)

    # 动态更新智能体位置
    for state in path:
        positions, time = state
        ax.cla()  # 清除当前图形

        # 重新绘制栅格和障碍物
        for x in range(width + 1):
            ax.axvline(x, color='black', linewidth=1)
        for y in range(height + 1):
            ax.axhline(y, color='black', linewidth=1)

        for x in range(width):
            for y in range(height):
                if map[y][x] == 1:
                    rect = plt.Rectangle((x, y), 1, 1, facecolor='gray')
                    ax.add_patch(rect)

        # 绘制智能体位置
        colors = ['red', 'blue', 'green', 'orange', 'purple']  # 每个智能体一种颜色
        for i, (x, y) in enumerate(positions):
            circle = plt.Circle((x + 0.5, y + 0.5), 0.3, color=colors[i % len(colors)])
            ax.add_patch(circle)

        # 更新标题以显示当前时间步
        ax.set_title(f'Time {time}')

        # 显示并暂停一段时间
        plt.pause(1)

    plt.show()


if __name__ == '__main__':
    # 示例用例
    map = [
        [1, 1, 0, 1, 1],
        [1, 1, 0, 1, 1],
        [0, 0, 0, 0, 0],
        [1, 1, 0, 1, 1],
        [1, 1, 0, 1, 1]
    ]

    start_positions = [(1, 2), (2, 1), (3, 2), (2, 3)]
    goal_positions = [(2, 1), (3, 2), (2, 3), (1, 2)]
    num_agents = 4

    # start_positions = [(1, 2), (2, 1), (3, 2), (2, 3), (2, 2), (0, 2), (2, 0), (4, 2)]
    # goal_positions = [(2, 1), (3, 2), (2, 3), (1, 2), (2, 2), (0, 2), (2, 0), (4, 2)]
    # num_agents = 8

    result = space_time_a_star(start_positions, goal_positions, map, num_agents)
    if result:
        print("Path found:")
        for state in result:
            positions, time = state
            print(f"Time {time}: {positions}")
            # 可视化路径
        visualize_path(result, map)
    else:
        print("No path found.")
